package com.lan.ai.mutimodelEmbedding;

import com.alibaba.fastjson.JSONObject;
import io.micrometer.observation.ObservationRegistry;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.MutationResult;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.dml.InsertParam;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.model.EmbeddingUtils;
import org.springframework.ai.vectorstore.MilvusVectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
import org.springframework.util.Assert;

import java.util.ArrayList;
import java.util.List;

/**
 * @author lanyanhua
 * @date 2025/1/21 16:14
 * @description
 */
public class MutimodeMilvusVectorStore extends MilvusVectorStore {
    public MutimodeMilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel, boolean initializeSchema) {
        this(milvusClient, embeddingModel, null, initializeSchema,
                new TokenCountBatchingStrategy());
    }

    public MutimodeMilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel, boolean initializeSchema, BatchingStrategy batchingStrategy) {
        this(milvusClient, embeddingModel, null, initializeSchema, batchingStrategy);
    }

    public MutimodeMilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel, MilvusVectorStoreConfig config, boolean initializeSchema, BatchingStrategy batchingStrategy) {
        this(milvusClient, embeddingModel, config, initializeSchema, batchingStrategy, ObservationRegistry.NOOP, null);
    }

    public MutimodeMilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel, MilvusVectorStoreConfig config, boolean initializeSchema, BatchingStrategy batchingStrategy, ObservationRegistry observationRegistry, VectorStoreObservationConvention customObservationConvention) {
        super(milvusClient, embeddingModel, MilvusVectorStore.MilvusVectorStoreConfig.builder()
                .withCollectionName(config.collectionName)
                .withDatabaseName(config.databaseName)
                .withIndexType(config.indexType)
                .withMetricType(config.metricType)
                .withIndexParameters(config.indexParameters)
                .withEmbeddingDimension(config.embeddingDimension)
                .build(), initializeSchema, batchingStrategy, observationRegistry, customObservationConvention);
        this.config = config;
        this.milvusClient = milvusClient;
        this.embeddingModel = embeddingModel;
        this.batchingStrategy = batchingStrategy;


    }

    private final MilvusServiceClient milvusClient;

    private final EmbeddingModel embeddingModel;

    private final MilvusVectorStoreConfig config;
    private final BatchingStrategy batchingStrategy;

    /**
     * Configuration for the Milvus vector store.
     */
    public record MilvusVectorStoreConfig(
            String databaseName,
            String collectionName,
            int embeddingDimension,
            IndexType indexType,
            MetricType metricType,
            String indexParameters) {

    }

    @Override
    public void doAdd(List<Document> documents) {

        Assert.notNull(documents, "Documents must not be null");

        List<String> docIdArray = new ArrayList<>();
        List<String> contentArray = new ArrayList<>();
        List<JSONObject> metadataArray = new ArrayList<>();
        List<List<Float>> embeddingArray = new ArrayList<>();

        // TODO: Need to customize how we pass the embedding options
        this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);

        for (Document document : documents) {
            docIdArray.add(document.getId());
            // Use a (future) DocumentTextLayoutFormatter instance to extract
            // the content used to compute the embeddings
            if(document instanceof MutimodeDocument) {
                //todo 处理多模态数据
                contentArray.add(((MutimodeDocument) document).sourceContent());
            }else {
                contentArray.add(document.getContent());
            }
            metadataArray.add(new JSONObject(document.getMetadata()));
            embeddingArray.add(EmbeddingUtils.toList(document.getEmbedding()));
        }

        List<InsertParam.Field> fields = new ArrayList<>();
        fields.add(new InsertParam.Field(DOC_ID_FIELD_NAME, docIdArray));
        fields.add(new InsertParam.Field(CONTENT_FIELD_NAME, contentArray));
        fields.add(new InsertParam.Field(METADATA_FIELD_NAME, metadataArray));
        fields.add(new InsertParam.Field(EMBEDDING_FIELD_NAME, embeddingArray));

        InsertParam insertParam = InsertParam.newBuilder()
                .withDatabaseName(this.config.databaseName)
                .withCollectionName(this.config.collectionName)
                .withFields(fields)
                .build();

        R<MutationResult> status = this.milvusClient.insert(insertParam);
        if (status.getException() != null) {
            throw new RuntimeException("Failed to insert:", status.getException());
        }
    }
}
